// 知识图谱系统

use std::collections::HashMap;
use crate::symbol_table::SymbolTable;
use crate::utils::tokenize;

// 实体类型枚举
#[derive(Debug, Clone, PartialEq)]
pub enum EntityType {
    Person,
    Location,
    Organization,
    Concept,
    Object,
    Event,
    Unknown,
}

// 关系类型枚举
#[derive(Debug, Clone, PartialEq)]
pub enum RelationType {
    IsA,          // 是一个
    PartOf,       // 是一部分
    LocatedIn,    // 位于
    HasProperty,  // 有属性
    CausedBy,     // 由...引起
    Causes,       // 导致
    RelatedTo,    // 相关
    Unknown,
}

// 知识实体
#[derive(Debug, Clone)]
pub struct KnowledgeEntity {
    pub id: usize,
    pub name: String,
    pub entity_type: EntityType,
    pub properties: HashMap<String, String>,
    pub last_accessed: u64,
}

// 知识关系
#[derive(Debug, Clone)]
pub struct KnowledgeRelation {
    pub id: usize,
    pub subject: usize,  // 主语实体ID
    pub object: usize,   // 宾语实体ID
    pub relation_type: RelationType,
    pub confidence: f32,
    pub last_accessed: u64,
}

// 知识图谱
#[derive(Debug)]
pub struct KnowledgeGraph {
    pub entities: HashMap<usize, KnowledgeEntity>,
    pub relations: HashMap<usize, KnowledgeRelation>,
    pub name_to_entity: HashMap<String, usize>,
    pub next_entity_id: usize,
    pub next_relation_id: usize,
    pub current_time: u64,
}

impl KnowledgeGraph {
    pub fn new() -> Self {
        Self {
            entities: HashMap::new(),
            relations: HashMap::new(),
            name_to_entity: HashMap::new(),
            next_entity_id: 0,
            next_relation_id: 0,
            current_time: 0,
        }
    }
    
    // 添加一些基本实体
    pub fn add_basic_entities(&mut self, symbol_table: &mut SymbolTable) {
        // 添加一些基本概念
        let basic_concepts = [
            ("人类", EntityType::Person),
            ("地方", EntityType::Location),
            ("组织", EntityType::Organization),
            ("概念", EntityType::Concept),
            ("物体", EntityType::Object),
            ("事件", EntityType::Event),
            ("时间", EntityType::Concept),
            ("空间", EntityType::Concept),
            ("原因", EntityType::Concept),
            ("结果", EntityType::Concept),
        ];
        
        for (name, entity_type) in basic_concepts {
            let name_id = symbol_table.intern(name);
            self.add_entity(name.to_string(), entity_type, name_id);
        }
    }
    
    // 添加实体
    pub fn add_entity(&mut self, name: String, entity_type: EntityType, _symbol_id: usize) -> usize {
        // 检查实体是否已存在
        if let Some(&id) = self.name_to_entity.get(&name) {
            // 更新最后访问时间
            if let Some(entity) = self.entities.get_mut(&id) {
                entity.last_accessed = self.current_time;
            }
            return id;
        }
        
        // 创建新实体
        let id = self.next_entity_id;
        self.next_entity_id += 1;
        
        let entity = KnowledgeEntity {
            id,
            name: name.clone(),
            entity_type,
            properties: HashMap::new(),
            last_accessed: self.current_time,
        };
        
        self.entities.insert(id, entity);
        self.name_to_entity.insert(name, id);
        
        id
    }
    
    // 添加关系
    pub fn add_relation(&mut self, subject_id: usize, object_id: usize, relation_type: RelationType, confidence: f32) -> usize {
        // 检查关系是否已存在
        for (_, relation) in &self.relations {
            if relation.subject == subject_id && 
               relation.object == object_id && 
               relation.relation_type == relation_type {
                // 更新置信度和最后访问时间
                let id = relation.id;
                if let Some(rel) = self.relations.get_mut(&id) {
                    rel.confidence = (rel.confidence + confidence) / 2.0; // 简单平均
                    rel.last_accessed = self.current_time;
                }
                return id;
            }
        }
        
        // 创建新关系
        let id = self.next_relation_id;
        self.next_relation_id += 1;
        
        let relation = KnowledgeRelation {
            id,
            subject: subject_id,
            object: object_id,
            relation_type,
            confidence,
            last_accessed: self.current_time,
        };
        
        self.relations.insert(id, relation);
        
        id
    }
    
    // 获取实体
    pub fn get_entity(&self, id: usize) -> Option<&KnowledgeEntity> {
        self.entities.get(&id)
    }
    
    // 根据名称获取实体
    pub fn get_entity_by_name(&self, name: &str) -> Option<&KnowledgeEntity> {
        if let Some(&id) = self.name_to_entity.get(name) {
            self.entities.get(&id)
        } else {
            None
        }
    }
    
    // 获取实体的关系
    pub fn get_entity_relations(&self, entity_id: usize) -> Vec<&KnowledgeRelation> {
        self.relations.values()
            .filter(|rel| rel.subject == entity_id || rel.object == entity_id)
            .collect()
    }
    
    // 查找相关实体
    pub fn find_related_entities(&self, entity_id: usize, relation_type: Option<RelationType>) -> Vec<usize> {
        let mut related = Vec::new();
        
        for relation in self.relations.values() {
            if relation.subject == entity_id {
                if let Some(rel_type) = &relation_type {
                    if relation.relation_type == *rel_type {
                        related.push(relation.object);
                    }
                } else {
                    related.push(relation.object);
                }
            } else if relation.object == entity_id {
                if let Some(rel_type) = &relation_type {
                    if relation.relation_type == *rel_type {
                        related.push(relation.subject);
                    }
                } else {
                    related.push(relation.subject);
                }
            }
        }
        
        related
    }
    
    // 更新时间
    pub fn update_time(&mut self) {
        self.current_time += 1;
    }
    
    // 从文本中提取实体
    pub fn extract_entities_from_text(&mut self, text: &str, symbol_table: &mut SymbolTable) -> Vec<usize> {
        let tokens = tokenize(text);
        let mut entities = Vec::new();
        
        // 简单的实体提取：大写开头的词作为潜在实体
        for token in tokens {
            if !token.is_empty() {
                let first_char = token.chars().next().unwrap();
                
                // 中文实体检测（简化）
                if crate::utils::is_chinese_text(&token) && token.len() <= 4 {
                    let token_id = symbol_table.intern(&token);
                    let entity_id = self.add_entity(token.clone(), EntityType::Unknown, token_id);
                    entities.push(entity_id);
                }
                // 英文实体检测（大写开头）
                else if first_char.is_uppercase() {
                    let token_id = symbol_table.intern(&token);
                    let entity_id = self.add_entity(token.clone(), EntityType::Unknown, token_id);
                    entities.push(entity_id);
                }
            }
        }
        
        entities
    }
    
    // 从对话中学习
    pub fn learn_from_dialogue(&mut self, input: &str, response: &str, symbol_table: &mut SymbolTable) {
        // 提取输入和回应中的实体
        let input_entities = self.extract_entities_from_text(input, symbol_table);
        let response_entities = self.extract_entities_from_text(response, symbol_table);
        
        // 建立实体之间的关系
        for &input_entity in &input_entities {
            for &response_entity in &response_entities {
                // 添加相关关系
                self.add_relation(input_entity, response_entity, RelationType::RelatedTo, 0.5);
            }
        }
        
        // 尝试识别因果关系
        self.extract_causal_relations(input, response, symbol_table);
    }
    
    // 提取因果关系
    fn extract_causal_relations(&mut self, input: &str, response: &str, symbol_table: &mut SymbolTable) {
        // 简单的因果关系提取：查找因果关键词
        let causal_keywords = ["因为", "由于", "导致", "引起", "所以", "因此", "because", "cause", "result", "therefore"];
        
        let input_tokens = tokenize(input);
        let _response_tokens = tokenize(response);
        
        // 查找输入中的因果关键词
        for (i, token) in input_tokens.iter().enumerate() {
            if causal_keywords.contains(&token.as_str()) {
                // 尝试识别原因和结果
                if i > 0 && i < input_tokens.len() - 1 {
                    let cause_token = &input_tokens[i-1];
                    let effect_token = &input_tokens[i+1];
                    
                    let cause_id = symbol_table.intern(cause_token);
                    let effect_id = symbol_table.intern(effect_token);
                    
                    let cause_entity = self.add_entity(cause_token.clone(), EntityType::Concept, cause_id);
                    let effect_entity = self.add_entity(effect_token.clone(), EntityType::Concept, effect_id);
                    
                    self.add_relation(cause_entity, effect_entity, RelationType::Causes, 0.7);
                }
            }
        }
    }
    
    // 查询知识图谱
    pub fn query(&self, query: &str, _symbol_table: &SymbolTable) -> Option<String> {
        let tokens = tokenize(query);
        
        // 简单的查询实现：查找实体
        for token in tokens {
            if let Some(entity) = self.get_entity_by_name(&token) {
                // 获取实体的关系
                let relations = self.get_entity_relations(entity.id);
                
                if !relations.is_empty() {
                    // 构建回应
                    let mut response = format!("{} 是 {}", entity.name, 
                        match entity.entity_type {
                            EntityType::Person => "一个人",
                            EntityType::Location => "一个地方",
                            EntityType::Organization => "一个组织",
                            EntityType::Concept => "一个概念",
                            EntityType::Object => "一个物体",
                            EntityType::Event => "一个事件",
                            EntityType::Unknown => "一个实体",
                        }
                    );
                    
                    // 添加关系信息
                    for relation in relations.iter().take(3) { // 限制关系数量
                        if let Some(subject) = self.get_entity(relation.subject) {
                            if let Some(object) = self.get_entity(relation.object) {
                                let relation_text = match relation.relation_type {
                                    RelationType::IsA => "是一个",
                                    RelationType::PartOf => "是...的一部分",
                                    RelationType::LocatedIn => "位于",
                                    RelationType::HasProperty => "有属性",
                                    RelationType::CausedBy => "由...引起",
                                    RelationType::Causes => "导致",
                                    RelationType::RelatedTo => "相关",
                                    RelationType::Unknown => "有关系",
                                };
                                
                                if relation.subject == entity.id {
                                    response += &format!(", {} {}", subject.name, relation_text);
                                } else {
                                    response += &format!(", {} {}", relation_text, object.name);
                                }
                            }
                        }
                    }
                    
                    return Some(response);
                }
            }
        }
        
        None
    }
}